package chat

type Hub struct {
	clients map[*Client]bool
	groups  map[int]map[*Client]bool

	broadcast      chan []byte
	broadcastGroup chan []byte

	register   chan *Client
	unregister chan *Client
}

func NewHub() *Hub {
	return &Hub{
		clients:        make(map[*Client]bool),
		groups:         make(map[int]map[*Client]bool),
		broadcast:      make(chan []byte),
		broadcastGroup: make(chan []byte),
		register:       make(chan *Client),
		unregister:     make(chan *Client),
	}
}

func (h *Hub) Run() {
	for {
		select {
		case client := <-h.register:
			h.Register(client)
		case client := <-h.unregister:
			h.Unregister(client)
		case message := <-h.broadcast:
			h.Broadcast(message)
		case message := <-h.broadcastGroup:
			h.Broadcast(message)
		}
	}
}

func (h *Hub) Register(c *Client) {
	h.clients[c] = true
	if _, ok := h.groups[c.group]; !ok {
		h.groups[c.group] = make(map[*Client]bool)
	}
	h.groups[c.group][c] = true
}

func (h *Hub) Unregister(c *Client) {
	if _, ok := h.clients[c]; ok {
		delete(h.clients, c)
	}
	if _, ok := h.groups[c.group][c]; ok {
		delete(h.groups[c.group], c)
	}
	close(c.send)
}

func (h *Hub) Broadcast(message []byte) {
	for client := range h.clients {
		select {
		case client.send <- message:
		default:
			// close(client.send)
			// delete(h.clients, client)
			h.Unregister(client)
		}
	}
}

func (h *Hub) BroadcastGroup(group int, message []byte) {
	for client := range h.groups[group] {
		select {
		case client.send <- message:
		default:
			// close(client.send)
			// delete(h.clients, client)
			h.Unregister(client)
		}
	}
}
